package cn.demoncat.util.lang.entity.data;

import cn.demoncat.util.lang.CollectionUtil;
import cn.demoncat.util.lang.IntegerUtil;
import cn.demoncat.util.lang.StringUtil;
import cn.demoncat.util.lang.constant.StatusConstant.Common;
import cn.demoncat.util.lang.constant.StringConstant;

import java.util.*;
import java.util.Map.Entry;
import java.util.function.Function;

/**
 * 赫夫曼树（最优二叉树）
 * 
 * @author 延晓磊
 *
 * @since 2020年7月21日
 */
public class HuffmanBinaryTree<T> extends BinaryTree<T> {

	/**
	 * 构建
	 * 
	 * @param arr			原数据
	 * @param weigher		原数据权重器（用于数据排序）
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	public static <T>HuffmanBinaryTree<T> build(T[] arr, Function<T,Integer> weigher) {
		// 将数据转换为Node列表
		List<BinaryTreeNode<T>> nodes = new ArrayList<>();
		for (T t : arr) {
			// 转换为Node，设置权值
			nodes.add(new BinaryTreeNode<>(t).setWeight(IntegerUtil.toEmpty(weigher.apply(t))));
		}
		// 构建赫夫曼树
		BinaryTreeNode<T> node;
		BinaryTreeNode<T> leftNode;
		BinaryTreeNode<T> rightNode;
		while (nodes.size() > 1) {
			// 排序：从小到大
			Collections.sort(nodes);
			// 取出WPL最小的叶子结点
			leftNode = nodes.get(0);
			rightNode = nodes.get(1);
			// 构建子树
			node = new BinaryTreeNode<T>().setWeight(leftNode.weight + rightNode.weight);
			node.left = leftNode;
			node.right = rightNode;
			// 删除处理过的结点
			nodes.remove(0);
			nodes.remove(0);
			// 添加子树
			nodes.add(node);
		}
		HuffmanBinaryTree<T> bt = new HuffmanBinaryTree<>();
		bt.setRoot(nodes.get(0));
		return bt;
	}

	/**
	 * 压缩 - 赫夫曼编码
	 * 
	 * @param data	原数据
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	public static ZipResult zip(byte[] data) {
		ZipResult result = new ZipResult();
		// 统计频率
		Map<Byte, Integer> counts = byteCount(data);
		// 构建赫夫曼树
		HuffmanBinaryTree<Byte> tree = build(CollectionUtil.toArray(counts.keySet(), new Byte[counts.size()]), counts::get);
		// 整合赫夫曼编码表
		result.setCodes(getCodes(tree.getRoot()));
		// 转换数据为编码
		StringBuilder code = new StringBuilder();
		for (byte b : data) {
			code.append(result.getCodes().get(b));
		}
		result.setLength(code.length());
		// 编码补0
		if (code.length() % 8 != 0) {
			for (int i = 0; i < code.length() % 8; i++) {
				code.append(StringConstant.N0);
			}
		}
		// 转换编码为byte[]
		result.setData(StringUtil.toBinaryBytes(code.toString()));
		return result;
	}
	
	/**
	 * 解压 - 赫夫曼解码
	 * 
	 * @param zip	编码结果
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	public static byte[] unzip(ZipResult zip) {
		return unzip(zip.getData(), zip.getCodes(), zip.getLength());
	}
	
	/**
	 * 解压 - 赫夫曼解码
	 * 
	 * @param data		编码数据
	 * @param codes		编码表
	 * @param length	编码长度
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	public static byte[] unzip(byte[] data, Map<Byte,String> codes, int length) {
		// 转换数据为编码
		String code = StringUtil.toBinaryStr(data);
		// 编码除0
		if (code.length() > length) {
			code = code.substring(0, length);
		}
		// 反向映射编码表
		Map<String,Byte> map = new HashMap<>(codes.size());
		for (Entry<Byte, String> entry : codes.entrySet()) {
			map.put(entry.getValue(), entry.getKey());
		}
		// 转换原数据
		return getBytes(code, map);
	}
	
	/**
	 * 转换原数据
	 * 
	 * @param code	编码
	 * @param map	反向编码表
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	private static byte[] getBytes(String code, Map<String,Byte> map) {
		List<Byte> list = new ArrayList<>();
		int count;
		Byte b;
		for (int i = 0; i < code.length(); i += count) {
			// 计数器
			count = 1;
			while (true) {
				// 读取编码，映射原数据
				if (i + count > code.length()) {
					break;
				}
				b = map.get(code.substring(i, i + count));
				if (b == null) {
					// 扩大读取范围
					count ++;
				}else {
					// 匹配成功
					list.add(b);
					break;
				}
			}
		}
		// 返回解码数据
		byte[] data = new byte[list.size()];
		for (int i = 0; i < data.length; i++) {
			data[i] = list.get(i);
		}
		return data;
	}
	
	/**
	 * 整合编码表
	 * 
	 * @param root	根结点
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	private static Map<Byte,String> getCodes(BinaryTreeNode<Byte> root){
		if (root == null) {
			return null;
		}
		Map<Byte,String> codeMap = new HashMap<>();
		// 左结点 0
		getCodes(root.left, Common.NO, null, codeMap);
		// 右结点 1
		getCodes(root.right, Common.YES, null, codeMap);
		return codeMap;
	}
	
	/**
	 * 整合编码表
	 * 
	 * @param node		当前结点
	 * @param code		路径编码：左结点0，右结点1
	 * @param codes		已拼接的路径编码
	 * @param codeMap	编码表
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	private static void getCodes(BinaryTreeNode<Byte> node, String code, StringBuilder codes, Map<Byte,String> codeMap){
		// 拼接结点路径
		codes = codes == null ? new StringBuilder(code) : new StringBuilder(codes).append(code);
		// 递归处理
		if (node != null) {
			if (node.data == null) {
				// 分枝结点，继续拼接
				getCodes(node.left, Common.NO, codes, codeMap);
				getCodes(node.right, Common.YES, codes, codeMap);
			}else {
				// 叶子结点，完成编码
				codeMap.put(node.data, codes.toString());
			}
		}
	}
	
	/**
	 * 字节计数
	 * 
	 * @param data	原数据
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	private static Map<Byte,Integer> byteCount(byte[] data){
		Map<Byte,Integer> counts = new HashMap<>();
		Integer count;
		for (byte b : data) {
			count = counts.get(b);
			if (count == null) {
				count = 1;
			}else {
				count ++;
			}
			counts.put(b, count);
		}
		return counts;
	}
	
	/**
	 * 压缩结果
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020年7月21日
	 */
	public static class ZipResult{
		
		// 转码数据：原数据 > 编码 > 转码数据 > 编码 > 原数据
		private byte[] data;
		// 编码表：编码时生成的编码表，解码时的编码表要与之一致
		private Map<Byte,String> codes;
		// 编码长度：编码时如果长度%8非0则右侧补0，解码时如果存在补0就相应去除，以此保证最后一个byte一致
		private int length;
		
		public byte[] getData() {
			return data;
		}
		public void setData(byte[] data) {
			this.data = data;
		}
		public Map<Byte, String> getCodes() {
			return codes;
		}
		public void setCodes(Map<Byte, String> codes) {
			this.codes = codes;
		}
		public int getLength() {
			return length;
		}
		public void setLength(int length) {
			this.length = length;
		}
	}
	
}
